'''
2024-4-3 Miao
单输入，单输出，分割模型，动态batch
'''
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import time
import cv2
import numpy as np
import torch.nn.functional as F
import torch
from PIL import Image
import torchvision.transforms as transforms
import os
          
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
             
def input_transform(image):
    image = image.astype(np.float32)[:, :, ::-1]
    image = image / 255.0
    image -= mean
    image /= std
    return image

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
f = open("*.engine", "rb")  # 添加自己的.engine文件路径，后缀不限制，也可以是.trt等
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING)) 
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()      

img = cv2.imread("*.png", cv2.IMREAD_COLOR)  # 添加自己的图片路径
sv_img = np.zeros_like(img).astype(np.uint8)
img = input_transform(img)
img = img.transpose((2, 0, 1)).copy()
img = img[np.newaxis, :, :, :].astype(np.float32)

context.set_binding_shape(0, tuple(img.shape))

h_input = cuda.pagelocked_empty(tuple(context.get_binding_shape(0)), dtype=np.float32) 
h_output = cuda.pagelocked_empty(tuple(context.get_binding_shape(1)), dtype=np.float32)

d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
bindings = [int(d_input), int(d_output)]

stream = cuda.Stream()


def predict(batch): 
    cuda.memcpy_htod_async(d_input, batch, stream)
    context.execute_async_v2(bindings, stream.handle)  
    cuda.memcpy_dtoh_async(output, d_output, stream)
    stream.synchronize()
    return output

t0 = time.time()
for i in range(1000):
    pred = predict(img)
t = time.time() - t0

print("Prediction cost {:.4f}s".format(t/1000))    # 输出每张图片的平均推理时间

d_input.free()
d_output.free()
   


